from torch import nn
from Encoding_labels import encoding


def Fast_RCNN_Loss (P, U, t_u, V, lamda=10, n_cls=1 ,n_reg = 1*9):
    """

    classification & regression loos 的计算
    :param P: 计算得到的label
    :param U: 原始的label
    :param t_u: 计算得到的坐标
    :param V: 原始的坐标
    :param lamda:
    :param n_cls:
    :param n_reg:
    :return:
    """
    cls_func = nn.CrossEntropyLoss()
    loc_func = nn.SmoothL1Loss()
    cls_loss = cls_func(P, U)
    loc_loss = loc_func(t_u, V)
    if U == encoding(None):  # 负例不计cls_loss
        IBI  = 0
    else:
        IBI = 1
    loss = (cls_loss / n_cls) + (lamda * IBI * loc_loss / n_reg)
    return loss